from typing import Any
import pymysql
from pymysql.connections import Connection

from dify_plugin import ToolProvider
from dify_plugin.errors.tool import ToolProviderCredentialValidationError


class ScarecrowMysqlProvider(ToolProvider):
    def _validate_credentials(self, credentials: dict[str, Any]) -> None:
        mysql_connection: Connection = None
        try:
            """
            验证提供的MySQL配置项是否有效。
            尝试使用提供的凭据建立数据库连接并执行测试查询。
            如果验证失败，将抛出 ToolProviderCredentialValidationError 异常。
            如果验证成功，将不会抛出异常。
            
            Args:
                credentials: 包含MySQL连接配置的字典，需包含以下键：
                    - mysql_host: 数据库主机地址
                    - mysql_port: 数据库端口号
                    - mysql_user_name: 数据库用户名
                    - mysql_user_pwd: 数据库密码
                    - mysql_database_name: 数据库名称
            """
            # 获取连接参数
            mysql_host = credentials["mysql_host"]
            mysql_port = int(credentials["mysql_port"])  # 确保端口号为整数
            mysql_user = credentials["mysql_user_name"]
            mysql_password = credentials["mysql_user_pwd"]
            mysql_database = credentials["mysql_database_name"]
            mysql_charset = "utf8mb4"  # 使用utf8mb4以支持完整的Unicode字符集
            
            # 建立数据库连接
            mysql_connection = pymysql.connect(
                host=mysql_host,
                port=mysql_port,
                user=mysql_user,        # 用户名
                password=mysql_password,  # 密码
                database=mysql_database,  # 数据库名
                charset=mysql_charset,    # 字符集
            )
            
            # 测试连接是否有效
            with mysql_connection.cursor() as cursor:
                cursor.execute("SELECT 1")
                cursor.fetchone()
                
        except KeyError as ke:
            raise ToolProviderCredentialValidationError(f"缺少必要的配置项: {str(ke)}")
        except ValueError as ve:
            raise ToolProviderCredentialValidationError(f"配置项格式错误: {str(ve)}")
        except Exception as e:
            raise ToolProviderCredentialValidationError(f"数据库连接验证失败: {str(e)}")
        finally:
            # 确保连接被正确关闭
            if mysql_connection:
                mysql_connection.close()
